#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : RAG-demo
# @File    : file_util.py
# @IDE     : PyCharm
# @Author  :ZH
# @Time    : 2025/6/30 15:00
import hashlib
from io import BytesIO
from typing import Optional, List

from fastapi import UploadFile

from util.async_http_client_util import async_http_client


async def calculate_md5(file: UploadFile) -> Optional[str]:
    """计算给定 UploadFile 对象的内容的 MD5 哈希值。

    Args:
        file (UploadFile): 从客户端上传的文件。

    Returns:
        Optional[str]: 文件内容的 MD5 哈希值，如果没有内容则返回 None。
    """
    # 初始化 MD5 对象
    hash_md5 = hashlib.md5()

    # 逐块读取文件内容并更新 MD5 对象
    while chunk := await file.read(8192):
        hash_md5.update(chunk)

    # 将文件指针重置到文件开头
    await file.seek(0)

    # 获取并返回 MD5 哈希值
    return hash_md5.hexdigest()


class ZhipuFileUtil:
    """
    智谱文件工具类，当前智谱的文件提取是免费的服务，所以暂用当前方案
    """

    def __init__(self, api_key: str = None, base_url: str = None):
        """

        :param api_key:
        :param base_url:
        """
        self.api_key = api_key
        self.base_url = base_url if base_url else 'https://open.bigmodel.cn/api/paas/v4/files'

    async def extract_file_content(self, file_object) -> str:
        """
        提取文件内容
        :param file_object:
        :return:
        """
        response = await async_http_client.post(
            url=self.base_url,
            files={
                'file': file_object,  # (file_name, BytesIO)
                'purpose': (None, 'file-extract')
            },
            headers=self._header()
        )
        file_id = response.get('id')
        content = await self.get_file_content(file_id=file_id)
        # 删除记录
        await self.delete_file(file_id=file_id)
        return content

    def _header(self) -> dict:
        return {"Authorization": f"Bearer {self.api_key}"}

    async def delete_file(self, file_id):
        await async_http_client.delete(
            url=f'https://open.bigmodel.cn/api/paas/v4/files/{file_id}',
            headers=self._header()
        )

    async def get_file_content(self, file_id):
        response = await async_http_client.get(
            url=f'https://open.bigmodel.cn/api/paas/v4/files/{file_id}/content',
            headers=self._header()
        )
        return response.get('content')


class VLLMFileUtil:
    """
    为方便修改，所以新建了另一个服务来处理识别问题
    """

    async def pdf_to_text(self, file_path: str) -> List[dict]:
        """
        使用VLLM服务将PDF文件转换为文本
        :param file_path: PDF文件的路径
        :return: 提取的文本内容
        """
        # 发送POST请求到VLLM服务
        with open(file_path, 'rb') as f:
            response = await async_http_client.post(
                url='http://127.0.0.1:8000/pdf/ocr',
                files={'file': (file_path, BytesIO(f.read()))},
            )
        # 获取每一页的内容
        return response.get('results', [{}])

    async def image_to_text(self, file_path: str) -> str:
        """
        使用VLLM服务将图片文件转换为文本
        :param file_path: 图片文件的路径
        :return: 提取的文本内容
        """
        # 发送POST请求到VLLM服务
        with open(file_path, 'rb') as f:
            response = await async_http_client.post(
                url='http://127.0.0.1:8000/image/ocr',
                files={'file': (file_path, BytesIO(f.read()))},
            )
        return response.get('content', "")


